{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b0532eb7",
   "metadata": {},
   "source": [
    "## RAG using Langchain with Neo4j Vector DB\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9ec40f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Check Dependencies\n",
    "#!pip3 install --upgrade --quiet langchain langchain-community langchain-openai langchain_mistralai neo4j==5.19.0 tiktoken tokenizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4108ab2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: langchain\n",
      "Version: 0.1.17\n",
      "Summary: Building applications with LLMs through composability\n",
      "Home-page: https://github.com/langchain-ai/langchain\n",
      "Author: \n",
      "Author-email: \n",
      "License: MIT\n",
      "Location: /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages\n",
      "Requires: aiohttp, dataclasses-json, jsonpatch, langchain-community, langchain-core, langchain-text-splitters, langsmith, numpy, pydantic, PyYAML, requests, SQLAlchemy, tenacity\n",
      "Required-by: \n",
      "---\n",
      "Name: neo4j\n",
      "Version: 5.19.0\n",
      "Summary: Neo4j Bolt driver for Python\n",
      "Home-page: \n",
      "Author: \n",
      "Author-email: \"Neo4j, Inc.\" <drivers@neo4j.com>\n",
      "License: Apache License, Version 2.0\n",
      "Location: /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages\n",
      "Requires: pytz\n",
      "Required-by: \n",
      "---\n",
      "Name: langchain-mistralai\n",
      "Version: 0.1.5\n",
      "Summary: An integration package connecting Mistral and LangChain\n",
      "Home-page: https://github.com/langchain-ai/langchain\n",
      "Author: \n",
      "Author-email: \n",
      "License: MIT\n",
      "Location: /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages\n",
      "Requires: httpx, httpx-sse, langchain-core, tokenizers\n",
      "Required-by: \n"
     ]
    }
   ],
   "source": [
    "!pip3 show langchain neo4j langchain_mistralai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57fe23fb",
   "metadata": {},
   "source": [
    "### Load environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1218f28",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NEO4J_URI = bolt://localhost:7687\n",
      "NEO4J_USERNAME = neo4j\n",
      "NEO4J_PASSWORD = neo4j123\n"
     ]
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "load_dotenv()\n",
    "\n",
    "print(\"NEO4J_URI = \" + os.getenv(\"NEO4J_URI\"))\n",
    "print(\"NEO4J_USERNAME = \" + os.getenv(\"NEO4J_USERNAME\"))\n",
    "print(\"NEO4J_PASSWORD = \" + os.getenv(\"NEO4J_PASSWORD\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f4595a8",
   "metadata": {},
   "source": [
    "### Create Neo4jGraph\n",
    "\n",
    "Make sure the Neo4j instance is running and the environment variables are set correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b87033b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.graphs import Neo4jGraph\n",
    "graph = Neo4jGraph()\n",
    "# graph = Neo4jGraph(url=\"bolt://localhost:7687\", username=\"neo4j\", password=\"neo4j123\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee99103f",
   "metadata": {},
   "source": [
    "### Check graph schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "39f09c26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node properties:\n",
      "Element {tags: STRING, name: STRING, type: STRING, embedding: LIST, source: STRING, parent: STRING, description: STRING}\n",
      "Relationship properties:\n",
      "Uses {consumer: STRING, provider: STRING, embedding: LIST, source: STRING, technology: STRING, description: STRING}\n",
      "The relationships:\n",
      "(:Element)-[:Uses]->(:Element)\n",
      "(:Element)-[:Contains]->(:Element)\n"
     ]
    }
   ],
   "source": [
    "print(graph.schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "287ca85b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# graph.structured_schema\n",
    "# graph.structured_schema['node_props']\n",
    "# graph.structured_schema['rel_props']\n",
    "# graph.structured_schema['relationships']\n",
    "# graph.structured_schema['metadata']['index']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b130e26b",
   "metadata": {},
   "source": [
    "### Cypher Query Examples\n",
    "\n",
    "We prepared some questions that will be passed by LLM.\n",
    "\n",
    "The first 4 questions are also provided with example Cypher query statements that should be used to retrieve relevant data from the graph database.\n",
    "\n",
    "They will be used to test the LLM's ability to generate Cypher query statements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "eebcc8ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_examples = [\n",
    "    {#1\n",
    "        \"question\": \"what software systems are used by customers?\",\n",
    "        \"query\": \"\"\"\n",
    "                MATCH (p:Element)-[r:Uses]->(ss:Element) \n",
    "                WHERE p.type=\"Person\" AND p.tags CONTAINS \"Customer\" AND ss.type=\"SoftwareSystem\" \n",
    "                RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\n",
    "                \"\"\",\n",
    "    },\n",
    "    {#2\n",
    "        \"question\": \"what software systems are used by customers to withdraw cash?\",\n",
    "        \"query\": \"\"\"\n",
    "                MATCH (p:Element)-[r:Uses]->(ss:Element) \n",
    "                WHERE p.type=\"Person\" AND p.tags CONTAINS \"Customer\" \n",
    "                    AND (r.description CONTAINS 'withdraw cash' OR (ss.type=\"SoftwareSystem\" AND ss.description CONTAINS 'withdraw cash'))\n",
    "                RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\n",
    "                \"\"\",\n",
    "    },\n",
    "    {#3\n",
    "        \"question\": \"what software systems are available?\",\n",
    "        \"query\": \"\"\"\n",
    "                MATCH (ss:Element) \n",
    "                WHERE ss.type='SoftwareSystem' \n",
    "                RETURN ss.tags + \":\" + ss.name as system, ss.description as description\n",
    "                \"\"\",\n",
    "    },\n",
    "    {#4\n",
    "        \"question\": \"what users are available?\",\n",
    "        \"query\": \"\"\"\n",
    "                MATCH (p:Element) \n",
    "                WHERE p.type='Person' \n",
    "                RETURN p.tags + \":\" + p.name as person, p.description as description\n",
    "                \"\"\",\n",
    "    },\n",
    "    {#5\n",
    "        \"question\": \"what software systems are used by staff?\",\n",
    "    },\n",
    "    {#6\n",
    "        \"question\": \"What software systems are used by Back Office Staff?\",\n",
    "    },\n",
    "    {#7\n",
    "        \"question\": \"what can provide a summary of a customer's bank accounts?\"\n",
    "    },\n",
    "    {#8\n",
    "        \"question\": \"what can be used to store customer information?\"\n",
    "    },\n",
    "    {#9\n",
    "        \"question\": \"what can be used by customers to view their banking information?\",\n",
    "    },\n",
    "    {#10\n",
    "        \"question\": \"what software system can be used to store customer information?\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb1cbaa4",
   "metadata": {},
   "source": [
    "### Simple query from graph database "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "399addd7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "      question: what software systems are used by customers?\n",
      "      query: \n",
      "                MATCH (p:Element)-[r:Uses]->(ss:Element) \n",
      "                WHERE p.type=\"Person\" AND p.tags CONTAINS \"Customer\" AND ss.type=\"SoftwareSystem\" \n",
      "                RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\n",
      "                \n",
      "      results: [{'context': 'Element,Person,Customer Withdraws cash using ATM', 'name': 'ATM', 'description': 'Allows customers to withdraw cash.'}, {'context': 'Element,Person,Customer Views account balances, and makes payments using Internet Banking System', 'name': 'Internet Banking System', 'description': 'Allows customers to view information about their bank accounts, and make payments.'}]\n",
      "      \n"
     ]
    }
   ],
   "source": [
    "print(f\"\"\"\n",
    "      question: {query_examples[0][\"question\"]}\n",
    "      query: {query_examples[0][\"query\"]}\n",
    "      results: {graph.query(query_examples[0][\"query\"])}\n",
    "      \"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc747d13",
   "metadata": {},
   "source": [
    "### Create LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dd6c9c3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MISTRAL_MODEL = open-mistral-7b\n"
     ]
    }
   ],
   "source": [
    "# from langchain_openai import ChatOpenAI\n",
    "# os.environ[\"OPENAI_API_KEY\"] = .......\n",
    "# llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "\n",
    "from langchain_mistralai.chat_models import ChatMistralAI\n",
    "llm = ChatMistralAI(mistral_api_key=os.getenv(\"MISTRAL_API_KEY\"), model=os.getenv(\"MISTRAL_MODEL\"), temperature=0)\n",
    "\n",
    "print(\"MISTRAL_MODEL = \" + os.getenv(\"MISTRAL_MODEL\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0900e165",
   "metadata": {},
   "source": [
    "### Create GraphCypherQAChain\n",
    "\n",
    "Please refer to API documentation for more details.\n",
    "\n",
    "https://api.python.langchain.com/en/latest/chains/langchain.chains.graph_qa.cypher.GraphCypherQAChain.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ddf43ecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import GraphCypherQAChain\n",
    "chain = GraphCypherQAChain.from_llm(graph=graph, llm=llm, verbose=True)\n",
    "\n",
    "# class GraphCypherQAChain(Chain):\n",
    "#     graph: GraphStore = Field(exclude=True)\n",
    "#     cypher_generation_chain: LLMChain\n",
    "#     qa_chain: LLMChain\n",
    "#     graph_schema: str\n",
    "#     input_key: str = \"query\"  #: :meta private:\n",
    "#     output_key: str = \"result\"  #: :meta private:\n",
    "#     top_k: int = 10\n",
    "#     \"\"\"Number of results to return from the query\"\"\"\n",
    "#     return_intermediate_steps: bool = False\n",
    "#     \"\"\"Whether or not to return the intermediate steps along with the final answer.\"\"\"\n",
    "#     return_direct: bool = False\n",
    "#     \"\"\"Whether or not to return the result of querying the graph directly.\"\"\"\n",
    "#     cypher_query_corrector: Optional[CypherQueryCorrector] = None\n",
    "#     \"\"\"Optional cypher validation tool\"\"\"\n",
    "\n",
    "# from_llm(\n",
    "#         cls,\n",
    "#         llm: Optional[BaseLanguageModel] = None,\n",
    "#         *,\n",
    "#         qa_prompt: Optional[BasePromptTemplate] = None,\n",
    "#         cypher_prompt: Optional[BasePromptTemplate] = None,\n",
    "#         cypher_llm: Optional[BaseLanguageModel] = None,\n",
    "#         qa_llm: Optional[BaseLanguageModel] = None,\n",
    "#         exclude_types: List[str] = [],\n",
    "#         include_types: List[str] = [],\n",
    "#         validate_cypher: bool = False,\n",
    "#         qa_llm_kwargs: Optional[Dict[str, Any]] = None,\n",
    "#         cypher_llm_kwargs: Optional[Dict[str, Any]] = None,\n",
    "#         **kwargs: Any,\n",
    "#     ) -> GraphCypherQAChain:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00c36309",
   "metadata": {},
   "source": [
    "### Check the default prompts\n",
    "\n",
    "CYPHER_GENERATION_PROMPT is used to generate the cypher statement.\n",
    "\n",
    "CYPHER_QA_PROMPT is used to generate the answer to questions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "588ab77d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CYPHER_GENERATION_PROMPT : \n",
      "input_variables=['question', 'schema'] template='Task:Generate Cypher statement to query a graph database.\\nInstructions:\\nUse only the provided relationship types and properties in the schema.\\nDo not use any other relationship types or properties that are not provided.\\nSchema:\\n{schema}\\nNote: Do not include any explanations or apologies in your responses.\\nDo not respond to any questions that might ask anything else than for you to construct a Cypher statement.\\nDo not include any text except the generated Cypher statement.\\n\\nThe question is:\\n{question}'\n",
      "CYPHER_QA_PROMPT : \n",
      "input_variables=['context', 'question'] template=\"You are an assistant that helps to form nice and human understandable answers.\\nThe information part contains the provided information that you must use to construct an answer.\\nThe provided information is authoritative, you must never doubt it or try to use your internal knowledge to correct it.\\nMake the answer sound as a response to the question. Do not mention that you based the result on the given information.\\nHere is an example:\\n\\nQuestion: Which managers own Neo4j stocks?\\nContext:[manager:CTL LLC, manager:JANE STREET GROUP LLC]\\nHelpful Answer: CTL LLC, JANE STREET GROUP LLC owns Neo4j stocks.\\n\\nFollow this example when generating answers.\\nIf the provided information is empty, say that you don't know the answer.\\nInformation:\\n{context}\\n\\nQuestion: {question}\\nHelpful Answer:\"\n"
     ]
    }
   ],
   "source": [
    "from langchain.chains.graph_qa.prompts import CYPHER_GENERATION_PROMPT, CYPHER_QA_PROMPT\n",
    "\n",
    "print(f\"CYPHER_GENERATION_PROMPT : \\n{CYPHER_GENERATION_PROMPT}\", flush=True)\n",
    "\n",
    "print(f\"CYPHER_QA_PROMPT : \\n{CYPHER_QA_PROMPT}\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2576bf47",
   "metadata": {},
   "source": [
    "### Ask GraphCypherQAChain\n",
    "\n",
    "The **question** must be passed as a dictionary with a key \"query\" and the value is the question.\n",
    "\n",
    "The QA chain will generate a cypher query statement using `CYPHER_GENERATION_PROMPT` and return the answer to the question.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "036634a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems are used by customers?\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mcypher\n",
      "MATCH (e1:Element)-[:Uses]->(e2:Element)\n",
      "WHERE e1.tags IN [\"software\"] AND e2.tags IN [\"customer\"]\n",
      "RETURN e1.name AS Software_Name, e2.name AS Customer_Name\n",
      "\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'what software systems are used by customers?',\n",
       " 'result': \"I don't have information on the specific software systems used by customers.\"}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[0][\"question\"]}\"\"\")\n",
    "\n",
    "chain.invoke({\"query\": query_examples[0][\"question\"]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fb35960",
   "metadata": {},
   "source": [
    "### Promt Refinement for GraphCypherQAChain\n",
    "\n",
    "The cypher query statement generated by the QA chain is not always correct.\n",
    "\n",
    "We can use the `MY_CYPHER_GENERATION_PROMPT` and `MY_CYPHER_QA_PROMPT` to refine the prompts and pass them to the QA chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "74c37a49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "# ## Refine the CYPHER_GENERATION_TEMPLATE by adding additional instructions below\n",
    "# 1. Do not include parameters in the MATCH statement.\n",
    "# 2. Only include parameters in the WHERE statement.\n",
    "# 3. Always use \"MATCH (p:Element)-[r:Uses]->(ss:Element)\" to lookup \"Uses\" relationships.\n",
    "# 4. Always use \"MATCH (p:Element)-[r:Contains]->(ss:Element)\" to lookup \"Contains\" relationships.\n",
    "\n",
    "MY_CYPHER_GENERATION_TEMPLATE_PREFIX = \"\"\"\n",
    "Task: Generate a syntactically correct cypher statement to query a graph database.\n",
    "Instructions:\n",
    "Use only the provided relationship types and properties in the schema.\n",
    "Do not use any other relationship types or properties that are not provided.\n",
    "Do not include parameters in the MATCH statement.\n",
    "Only include parameters in the WHERE statement.\n",
    "Always use \"MATCH (p:Element)-[r:Uses]->(ss:Element)\" to lookup \"Uses\" relationships.\n",
    "Always use \"MATCH (p:Element)-[r:Contains]->(ss:Element)\" to lookup \"Contains\" relationships.\n",
    "\n",
    "Schema:\n",
    "{schema}\n",
    "\n",
    "Note: Do not include any explanations or apologies in your responses.\n",
    "Do not respond to any questions that might ask anything else than for you to construct a Cypher statement.\n",
    "Do not include any text except the generated Cypher statement.\n",
    "\n",
    "Here are a few examples of generated Cypher statements for particular questions:\n",
    "\"\"\"\n",
    "\n",
    "############################################\n",
    "##### Prompt refinement with examples ######\n",
    "############################################\n",
    "MY_CYPHER_GENERATION_TEMPLATE_EXAMPLES = f\"\"\"\n",
    "# {query_examples[0][\"question\"]} \n",
    "{query_examples[0][\"query\"]}\n",
    "# {query_examples[1][\"question\"]} \n",
    "{query_examples[1][\"query\"]}\n",
    "# {query_examples[2][\"question\"]} \n",
    "{query_examples[2][\"query\"]}\n",
    "\"\"\"\n",
    "\n",
    "MY_CYPHER_GENERATION_TEMPLATE_SUFFIX = \"\"\"   \n",
    "User input: {question}\n",
    "Cypher query:\n",
    "\"\"\"\n",
    "\n",
    "MY_CYPHER_GENERATION_TEMPLATE = MY_CYPHER_GENERATION_TEMPLATE_PREFIX + MY_CYPHER_GENERATION_TEMPLATE_EXAMPLES + MY_CYPHER_GENERATION_TEMPLATE_SUFFIX\n",
    "\n",
    "MY_CYPHER_GENERATION_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"schema\", \"question\"], template=MY_CYPHER_GENERATION_TEMPLATE\n",
    ")\n",
    "\n",
    "\n",
    "MY_CYPHER_QA_TEMPLATE = \"\"\"\n",
    "You are an assistant that helps to form nice and human understandable answers.\n",
    "The context part contains the provided information that you must use to construct an answer.\n",
    "The provided information is authoritative, you must never doubt it or try to use your internal knowledge to correct it.\n",
    "Make the answer sound as a response to the question. \n",
    "Do not mention that you based the result on the given context.\n",
    "Do not make up anything which does not exist in the provided context.\n",
    "\n",
    "Here is an example:\n",
    "Question: What software systems are used by customers?\n",
    "Context: Internet Banking System\n",
    "Helpful Answer: Internet Banking System.\n",
    "\n",
    "Follow this example when generating answers. If the provided context is empty, say that you don't know the answer.\n",
    "\n",
    "Context: {context}\n",
    "\n",
    "Question: {question}\n",
    "Helpful Answer:\n",
    "\"\"\"\n",
    "\n",
    "MY_CYPHER_QA_PROMPT = PromptTemplate(\n",
    "    input_variables=['context', 'question'], template=MY_CYPHER_QA_TEMPLATE\n",
    ")\n",
    "# Pass the refined prompts to the QA chain\n",
    "chain = GraphCypherQAChain.from_llm(llm=llm, graph=graph, cypher_prompt=MY_CYPHER_GENERATION_PROMPT,qa_prompt=MY_CYPHER_QA_PROMPT,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "68929886",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems are used by customers?\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Element)-[r:Uses]->(ss:Element)\n",
      "WHERE p.type=\"Person\" AND p.tags CONTAINS \"Customer\" AND ss.type=\"SoftwareSystem\"\n",
      "RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'context': 'Element,Person,Customer Withdraws cash using ATM', 'name': 'ATM', 'description': 'Allows customers to withdraw cash.'}, {'context': 'Element,Person,Customer Views account balances, and makes payments using Internet Banking System', 'name': 'Internet Banking System', 'description': 'Allows customers to view information about their bank accounts, and make payments.'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'what software systems are used by customers?',\n",
       " 'result': 'Customers use the Internet Banking System and ATMs. The Internet Banking System allows them to view account balances and make payments, while ATMs enable them to withdraw cash.'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[0][\"question\"]}\"\"\")\n",
    "\n",
    "chain.invoke({\"query\": query_examples[0][\"question\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e9bc8929",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems are available?\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (ss:Element)\n",
      "WHERE ss.type='SoftwareSystem'\n",
      "RETURN ss.tags + \":\" + ss.name as system, ss.description as description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'system': 'Element,Software System,Existing System:Mainframe Banking System', 'description': 'Stores all of the core banking information about customers, accounts, transactions, etc.'}, {'system': 'Element,Software System,Existing System:E-mail System', 'description': 'The internal Microsoft Exchange e-mail system.'}, {'system': 'Element,Software System,Existing System:ATM', 'description': 'Allows customers to withdraw cash.'}, {'system': 'Element,Software System:Internet Banking System', 'description': 'Allows customers to view information about their bank accounts, and make payments.'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'what software systems are available?',\n",
       " 'result': 'The software systems available include the Mainframe Banking System, the Microsoft Exchange e-mail system, ATMs, and the Internet Banking System.'}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[2][\"question\"]}\"\"\")\n",
    "\n",
    "chain.invoke({\"query\": query_examples[2][\"question\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "934c7704",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems are used by staff?\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Element)-[r:Uses]->(ss:Element)\n",
      "WHERE p.type=\"Person\" AND p.tags CONTAINS \"Staff\" AND ss.type=\"SoftwareSystem\"\n",
      "RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'context': 'Element,Person,Bank Staff Uses Mainframe Banking System', 'name': 'Mainframe Banking System', 'description': 'Stores all of the core banking information about customers, accounts, transactions, etc.'}, {'context': 'Element,Person,Bank Staff Uses Mainframe Banking System', 'name': 'Mainframe Banking System', 'description': 'Stores all of the core banking information about customers, accounts, transactions, etc.'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'what software systems are used by staff?',\n",
       " 'result': 'Mainframe Banking System is used by bank staff.'}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[4][\"question\"]}\"\"\")\n",
    "\n",
    "chain.invoke({\"query\": query_examples[4][\"question\"]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "214f2f09",
   "metadata": {},
   "source": [
    "### Few-shot examples\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "abeedf5c-8735-4419-9a08-324dd5554dc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate\n",
    "\n",
    "example_prompt = PromptTemplate.from_template(\n",
    "    \"User input: {question}\\nCypher query: {query}\"\n",
    ")\n",
    "\n",
    "fewShotPrompt = FewShotPromptTemplate(\n",
    "    examples=query_examples[:4],\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=MY_CYPHER_GENERATION_TEMPLATE_PREFIX,\n",
    "    suffix=MY_CYPHER_GENERATION_TEMPLATE_SUFFIX,\n",
    "    input_variables=[\"question\", \"schema\"],\n",
    ")\n",
    "\n",
    "chain = GraphCypherQAChain.from_llm(graph=graph, llm=llm, cypher_prompt=fewShotPrompt, qa_prompt=MY_CYPHER_QA_PROMPT, verbose=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2ae5b263",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(fewShotPrompt.format(question=query_examples[6][\"question\"], schema=graph.structured_schema))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8be7b78",
   "metadata": {},
   "source": [
    "### Ask a predefined question\n",
    "\n",
    "Ask a predefined question and the predefined query should be used to generate the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f00366df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems are used by customers?\n",
      "query: \n",
      "                MATCH (p:Element)-[r:Uses]->(ss:Element) \n",
      "                WHERE p.type=\"Person\" AND p.tags CONTAINS \"Customer\" AND ss.type=\"SoftwareSystem\" \n",
      "                RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\n",
      "                \n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Element)-[r:Uses]->(ss:Element)\n",
      " WHERE p.type=\"Person\" AND p.tags CONTAINS \"Customer\" AND ss.type=\"SoftwareSystem\"\n",
      " RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'context': 'Element,Person,Customer Withdraws cash using ATM', 'name': 'ATM', 'description': 'Allows customers to withdraw cash.'}, {'context': 'Element,Person,Customer Views account balances, and makes payments using Internet Banking System', 'name': 'Internet Banking System', 'description': 'Allows customers to view information about their bank accounts, and make payments.'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'what software systems are used by customers?',\n",
       " 'result': 'Customers use the Internet Banking System and ATMs. The Internet Banking System allows them to view account balances and make payments, while ATMs enable them to withdraw cash.'}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[0][\"question\"]}\\nquery: {query_examples[0][\"query\"]}\"\"\")\n",
    "\n",
    "chain.invoke(query_examples[0][\"question\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "090f50ef",
   "metadata": {},
   "source": [
    "### Ask a question that is not in the predefined list\n",
    "\n",
    "The QA Chain should be able to find the similar question in the predefined list and use the corresponding query to generate the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f2da7faf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems are used by staff?\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Element)-[r:Uses]->(ss:Element)\n",
      " WHERE p.type=\"Person\" AND p.tags CONTAINS \"Staff\" AND ss.type=\"SoftwareSystem\"\n",
      " RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'context': 'Element,Person,Bank Staff Uses Mainframe Banking System', 'name': 'Mainframe Banking System', 'description': 'Stores all of the core banking information about customers, accounts, transactions, etc.'}, {'context': 'Element,Person,Bank Staff Uses Mainframe Banking System', 'name': 'Mainframe Banking System', 'description': 'Stores all of the core banking information about customers, accounts, transactions, etc.'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'what software systems are used by staff?',\n",
       " 'result': 'Mainframe Banking System is used by bank staff.'}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[4][\"question\"]}\"\"\")\n",
    "\n",
    "chain.invoke(query_examples[4][\"question\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3c0a705d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: What software systems are used by Back Office Staff?\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphCypherQAChain chain...\u001b[0m\n",
      "Generated Cypher:\n",
      "\u001b[32;1m\u001b[1;3mMATCH (p:Element)-[r:Uses]->(ss:Element)\n",
      " WHERE p.type=\"Person\" AND p.tags CONTAINS \"Back Office Staff\" AND ss.type=\"SoftwareSystem\"\n",
      " RETURN p.tags + \" \" + r.description + \" \" + ss.name as context, ss.name as name, ss.description as description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'query': 'What software systems are used by Back Office Staff?',\n",
       " 'result': 'I cannot provide a definitive answer without specific information about the software systems used by Back Office Staff.'}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[5][\"question\"]}\"\"\")\n",
    "\n",
    "chain.invoke(query_examples[5][\"question\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4500a922",
   "metadata": {},
   "source": [
    "### Semantic Search\n",
    "\n",
    "The cypher statements generated for the above examples are using simple matching against keywords, e.g. `p.tags contains 'Bank Staff'` \n",
    "\n",
    "if we ask the question in another way, QA Chain will not be able to answer correctly, it should use semnatic search instead to resolve the problem. \n",
    "\n",
    "https://python.langchain.com/docs/use_cases/graph/semantic/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "cd427404",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(f\"\"\"question: {query_examples[7][\"question\"]}\"\"\")\n",
    "\n",
    "# chain.invoke(query_examples[7][\"question\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d5eeb77a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chain.invoke(\"what software systems can be used by customer to view their banking information?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0797c404",
   "metadata": {},
   "source": [
    "##### Create Neo4j Vector Store from an existing graph database\n",
    "\n",
    "Create embedding vector for the SoftwareSystem node from existing neo4j graph database. A new **embedding** property will be added to the SoftwareSystem node.\n",
    "\n",
    "- [MistralAIEmbeddings](https://python.langchain.com/docs/integrations/text_embedding/mistralai/)\n",
    "\n",
    "- [Neo4j Vector](https://python.langchain.com/docs/integrations/vectorstores/neo4jvector/)\n",
    "\n",
    "- [Create Neo4jVector from an existing graph](https://api.python.langchain.com/en/latest/_modules/langchain_community/vectorstores/neo4j_vector.html#Neo4jVector.from_existing_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "1d0d2980",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/langchain_mistralai/embeddings.py:105: UserWarning: Could not download mistral tokenizer from Huggingface for calculating batch sizes. Set a Huggingface token via the HF_TOKEN environment variable to download the real tokenizer. Falling back to a dummy tokenizer that uses `len()`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.vectorstores import Neo4jVector\n",
    "from langchain_mistralai import MistralAIEmbeddings\n",
    "\n",
    "retrieval_query = \"\"\"\n",
    "RETURN \"name: \"+node.name+\"\\ntype: \"+node.type+\"\\ndescription: \"+node.description AS text, score, {tags: node.tags, name: node.name, type: node.type, parent: node.parent, source: node.source} AS metadata\n",
    "\"\"\"\n",
    "\n",
    "# Initialize and return a Neo4jVector instance from an existing graph.\n",
    "# This method initializes a Neo4jVector instance using the provided parameters and the existing graph. \n",
    "# It validates the existence of the indices and creates new ones if they don't exist.\n",
    "vectorestore = Neo4jVector.from_existing_graph(\n",
    "    embedding=MistralAIEmbeddings(mistral_api_key=os.getenv(\"MISTRAL_API_KEY\")),\n",
    "    node_label=\"Element\",\n",
    "    embedding_node_property=\"embedding\",\n",
    "    text_node_properties=[\"description\"], #combime properties to create a single text property\n",
    "    index_name=\"Element_index\",\n",
    "    retrieval_query=retrieval_query, #default retrival query will contain all the text node properties\n",
    "    url=os.getenv(\"NEO4J_URI\"), username=os.getenv(\"NEO4J_USERNAME\"), password=os.getenv(\"NEO4J_PASSWORD\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "86082cdc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what can be used by customers to view their banking information?\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Document(page_content='name: Internet Banking System\\ntype: SoftwareSystem\\ndescription: Allows customers to view information about their bank accounts, and make payments.', metadata={'tags': 'Element,Software System', 'source': 'Big Bank plc', 'name': 'Internet Banking System', 'parent': '', 'type': 'SoftwareSystem'})]"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[8][\"question\"]}\"\"\")\n",
    "\n",
    "vectorestore.similarity_search(query_examples[8][\"question\"], k=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "9845a067",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='name: Mainframe Banking System\\ntype: SoftwareSystem\\ndescription: Stores all of the core banking information about customers, accounts, transactions, etc.', metadata={'tags': 'Element,Software System,Existing System', 'source': 'Big Bank plc', 'name': 'Mainframe Banking System', 'parent': '', 'type': 'SoftwareSystem'})]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = vectorestore.similarity_search(\"Mainframe Banking System\", k=1)\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "476f7896",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: name: Mainframe Banking System\n",
      "type: SoftwareSystem\n",
      "description: Stores all of the core banking information about customers, accounts, transactions, etc.\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, format_document\n",
    "\n",
    "# DEFAULT_DOCUMENT_PROMPT = ChatPromptTemplate.from_template(template=\"\\nname: {name}\\ntype: {type}{page_content}\")\n",
    "DEFAULT_DOCUMENT_PROMPT = ChatPromptTemplate.from_template(template=\"{page_content}\")\n",
    "def _extract_context_documents(\n",
    "    docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator=\"\\n\\n\"\n",
    "):\n",
    "    doc_strings = [format_document(doc, document_prompt) for doc in docs]\n",
    "    context_content = document_separator.join(doc_strings)\n",
    "    \n",
    "    return context_content\n",
    "\n",
    "print(_extract_context_documents(results))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "06bb2b8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what can be used by customers to view their banking information?\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[(Document(page_content='name: Internet Banking System\\ntype: SoftwareSystem\\ndescription: Allows customers to view information about their bank accounts, and make payments.', metadata={'tags': 'Element,Software System', 'source': 'Big Bank plc', 'name': 'Internet Banking System', 'parent': '', 'type': 'SoftwareSystem'}),\n",
       "  0.8344582915306091)]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[8][\"question\"]}\"\"\")\n",
    "\n",
    "results = vectorestore.similarity_search_with_score(query_examples[8][\"question\"], k=1)\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa37226b",
   "metadata": {},
   "source": [
    "#### Create Neo4j Vector Store from an existing embedding index\n",
    "\n",
    "- [MistralAIEmbeddings](https://python.langchain.com/docs/integrations/text_embedding/mistralai/)\n",
    "\n",
    "- [Neo4j Vector](https://python.langchain.com/docs/integrations/vectorstores/neo4jvector/)\n",
    "\n",
    "- [Create Neo4jVector from an existing index](https://api.python.langchain.com/en/latest/_modules/langchain_community/vectorstores/neo4j_vector.html#Neo4jVector.from_existing_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "42508a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Neo4jVector instance from an existing index.\n",
    "from langchain_community.vectorstores import Neo4jVector\n",
    "from langchain_mistralai import MistralAIEmbeddings\n",
    "\n",
    "embedding=MistralAIEmbeddings(mistral_api_key=os.getenv(\"MISTRAL_API_KEY\"))\n",
    "\n",
    "retrieval_query = \"\"\"\n",
    "RETURN \"\\n\\nname: \"+node.name+\"\\ntype: \"+node.type+\"\\ndescription: \"+node.description AS text, score, \n",
    "{tags: node.tags, name: node.name, type: node.type, parent: node.parent, source: node.source} AS metadata\n",
    "\"\"\"\n",
    "\n",
    "# Get instance of an existing Neo4j vector index. This method will return the instance of the store without inserting any new embeddings.\n",
    "vectorestore = Neo4jVector.from_existing_index(\n",
    "    embedding=embedding,\n",
    "    url=os.getenv(\"NEO4J_URI\"), username=os.getenv(\"NEO4J_USERNAME\"), password=os.getenv(\"NEO4J_PASSWORD\"),\n",
    "    index_name=\"Element_index\",\n",
    "    retrieval_query=retrieval_query,\n",
    "    embedding_node_property=\"embedding\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "08add5c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding_dimension: 1024\n",
      "index_type: NODE\n"
     ]
    }
   ],
   "source": [
    "\n",
    "embedding_dimension, index_type = vectorestore.retrieve_existing_index()\n",
    "\n",
    "print(f\"embedding_dimension: {embedding_dimension}\\nindex_type: {index_type}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "98de365e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what software systems can be used to store customer information?\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Document(page_content='\\n\\nname: Mainframe Banking System\\ntype: SoftwareSystem\\ndescription: Stores all of the core banking information about customers, accounts, transactions, etc.', metadata={'tags': 'Element,Software System,Existing System', 'source': 'Big Bank plc', 'name': 'Mainframe Banking System', 'parent': '', 'type': 'SoftwareSystem'})]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question = \"what software systems can be used to store customer information?\"\n",
    "print(f\"\"\"question: {question}\"\"\")\n",
    "results = vectorestore.similarity_search(query=question, k=1)\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "5fa48cc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: what can be used to store customer information?\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Document(page_content='\\n\\nname: Database\\ntype: Container\\ndescription: Stores user registration information, hashed authentication credentials, access logs, etc.', metadata={'tags': 'Element,Container,Database', 'source': 'Big Bank plc', 'name': 'Database', 'parent': 'Internet Banking System', 'type': 'Container'})]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"\"\"question: {query_examples[7][\"question\"]}\"\"\")\n",
    "results = vectorestore.similarity_search(query_examples[7][\"question\"], k=1)\n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "71354274",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='\\n\\nname: Internet Banking System\\ntype: SoftwareSystem\\ndescription: Allows customers to view information about their bank accounts, and make payments.', metadata={'tags': 'Element,Software System', 'source': 'Big Bank plc', 'name': 'Internet Banking System', 'parent': '', 'type': 'SoftwareSystem'})]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = vectorestore.similarity_search(\"what software systems can be used to withdraw cash?\", k=1)\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61870e53",
   "metadata": {},
   "source": [
    "### RAG - RetrievalQAWithSourcesChain using Element_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9c4d9f9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vectorestore.as_retriever()\n",
    "\n",
    "# from langchain_core.runnables import RunnableParallel, RunnablePassthrough\n",
    "# retrieval = RunnableParallel(\n",
    "#     {\"question\": RunnablePassthrough() | retriever | _extract_context_documents}\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "15f37d24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQAWithSourcesChain\n",
    "chain = RetrievalQAWithSourcesChain.from_chain_type(\n",
    "    llm, chain_type=\"stuff\", retriever=retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "4f2e019e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'The following software systems can be used to store customer information: a) Mainframe Banking System, b) Internet Banking System, and c) E-mail System.\\n\\n',\n",
       " 'sources': 'Big Bank plc (28-pl, 30-pl)'}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\n",
    "    {\"question\": \"what software systems can be used to store customer information?\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "22ddcd55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'The Mainframe Banking System, Internet Banking System, and ATM software systems can be used to store and manage customer banking information.\\n\\n',\n",
       " 'sources': 'Big Bank plc (28-pl, 30-pl, and multiple other instances)'}"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\n",
    "    {\"question\": \"what software systems can be used to store customer banking information?\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "bbf2fb5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'Customer information can be stored in a database, which is a container that stores user registration information, hashed authentication credentials, and access logs, among other things (Big Bank plc, ',\n",
       " 'sources': '1). The core banking information about customers, accounts, and transactions can also be stored in a mainframe banking system, which is a software system (Big Bank plc,'}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\n",
    "    {\"question\": \"what can be used to store customer information?\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "8104fb1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'The software systems that can be used to withdraw cash are ATM systems.\\n\\n',\n",
       " 'sources': 'Big Bank plc.'}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\n",
    "    {\"question\": \"what software systems can be used to withdraw cash?\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "655c49a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': \"what can provide a summary of a customer's bank accounts\",\n",
       " 'answer': 'The \"Accounts Summary Controller\" component in the \"Big Bank plc\" system provides a summary of a customer\\'s bank accounts.\\n\\n',\n",
       " 'sources': 'Big Bank plc (multiple sources)'}"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "chain.invoke(\n",
    "    {\"question\": \"what can provide a summary of a customer's bank accounts\"},\n",
    "    return_intermediate_steps=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5e791215",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(Document(page_content='\\n\\nname: Accounts Summary Controller\\ntype: Component\\ndescription: Provides customers with a summary of their bank accounts.', metadata={'tags': 'Element,Component', 'source': 'Big Bank plc', 'name': 'Accounts Summary Controller', 'parent': 'API Application', 'type': 'Component'}),\n",
       "  0.8787081241607666),\n",
       " (Document(page_content='\\n\\nname: Personal Banking Customer\\ntype: Person\\ndescription: A customer of the bank, with personal bank accounts.', metadata={'tags': 'Element,Person,Customer', 'source': 'Big Bank plc', 'name': 'Personal Banking Customer', 'parent': '', 'type': 'Person'}),\n",
       "  0.8426058292388916),\n",
       " (Document(page_content='\\n\\nname: Internet Banking System\\ntype: SoftwareSystem\\ndescription: Allows customers to view information about their bank accounts, and make payments.', metadata={'tags': 'Element,Software System', 'source': 'Big Bank plc', 'name': 'Internet Banking System', 'parent': '', 'type': 'SoftwareSystem'}),\n",
       "  0.8267369866371155),\n",
       " (Document(page_content='\\n\\nname: Mainframe Banking System\\ntype: SoftwareSystem\\ndescription: Stores all of the core banking information about customers, accounts, transactions, etc.', metadata={'tags': 'Element,Software System,Existing System', 'source': 'Big Bank plc', 'name': 'Mainframe Banking System', 'parent': '', 'type': 'SoftwareSystem'}),\n",
       "  0.8265427350997925)]"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = vectorestore.similarity_search_with_score(\"what can provide a summary of a customer's bank accounts?\")\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6439ee98",
   "metadata": {},
   "source": [
    "### Relatioship Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "652f2400",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First create relationship index and relationship data\n",
    "\n",
    "# Create a vector from an existing relationship index\n",
    "\n",
    "retrieval_query = \"\"\"\n",
    "RETURN \"'\"+relationship.consumer + \"' \" + relationship.description + \" '\" +relationship.provider + \"'\" AS text, score, \n",
    "{technology: relationship.technology, source: relationship.source, consumer: relationship.consumer, provider: relationship.provider} AS metadata\n",
    "\"\"\"\n",
    "\n",
    "relationship_vector = Neo4jVector.from_existing_relationship_index(\n",
    "    MistralAIEmbeddings(mistral_api_key=os.getenv(\"MISTRAL_API_KEY\")),\n",
    "    url=os.getenv(\"NEO4J_URI\"), username=os.getenv(\"NEO4J_USERNAME\"), password=os.getenv(\"NEO4J_PASSWORD\"),\n",
    "    index_name=\"Relationship_index\",\n",
    "    # text_node_property=\"text\", # the text property used as the page_content of Document\n",
    "    retrieval_query=retrieval_query, # the retrieval_query used as the page_content of Document\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "b394bb3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content=\"'Customer Service Staff' Uses 'Mainframe Banking System'\", metadata={'technology': '', 'source': 'Big Bank plc', 'provider': 'Mainframe Banking System', 'consumer': 'Customer Service Staff'})]"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relationship_vector.similarity_search(\"What system is used by Service Staff\", k=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "e1d8d0fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(Document(page_content=\"'Customer Service Staff' Uses 'Mainframe Banking System'\", metadata={'technology': '', 'source': 'Big Bank plc', 'provider': 'Mainframe Banking System', 'consumer': 'Customer Service Staff'}),\n",
       "  0.8858469724655151),\n",
       " (Document(page_content=\"'Back Office Staff' Uses 'Mainframe Banking System'\", metadata={'technology': '', 'source': 'Big Bank plc', 'provider': 'Mainframe Banking System', 'consumer': 'Back Office Staff'}),\n",
       "  0.8771942853927612),\n",
       " (Document(page_content=\"'Internet Banking System' Gets account information from, and makes payments using 'Mainframe Banking System'\", metadata={'technology': '', 'source': 'Big Bank plc', 'provider': 'Mainframe Banking System', 'consumer': 'Internet Banking System'}),\n",
       "  0.8604210019111633),\n",
       " (Document(page_content=\"'Internet Banking System' Sends e-mail using 'E-mail System'\", metadata={'technology': '', 'source': 'Big Bank plc', 'provider': 'E-mail System', 'consumer': 'Internet Banking System'}),\n",
       "  0.8592248558998108)]"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relationship_vector.similarity_search_with_score(\"What system is used by Service Staff\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "ea03dd7f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(Document(page_content='\\n\\nname: Customer Service Staff\\ntype: Person\\ndescription: Customer service staff within the bank.', metadata={'tags': 'Element,Person,Bank Staff', 'source': 'Big Bank plc', 'name': 'Customer Service Staff', 'parent': '', 'type': 'Person'}),\n",
       "  0.8767651319503784),\n",
       " (Document(page_content='\\n\\nname: E-mail System\\ntype: SoftwareSystem\\ndescription: The internal Microsoft Exchange e-mail system.', metadata={'tags': 'Element,Software System,Existing System', 'source': 'Big Bank plc', 'name': 'E-mail System', 'parent': '', 'type': 'SoftwareSystem'}),\n",
       "  0.858390748500824),\n",
       " (Document(page_content='\\n\\nname: Back Office Staff\\ntype: Person\\ndescription: Administration and support staff within the bank.', metadata={'tags': 'Element,Person,Bank Staff', 'source': 'Big Bank plc', 'name': 'Back Office Staff', 'parent': '', 'type': 'Person'}),\n",
       "  0.8562641143798828),\n",
       " (Document(page_content='\\n\\nname: Internet Banking System\\ntype: SoftwareSystem\\ndescription: Allows customers to view information about their bank accounts, and make payments.', metadata={'tags': 'Element,Software System', 'source': 'Big Bank plc', 'name': 'Internet Banking System', 'parent': '', 'type': 'SoftwareSystem'}),\n",
       "  0.850509762763977)]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vectorestore.similarity_search_with_score(\"What system is used by Service Staff\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c03a1a8",
   "metadata": {},
   "source": [
    "#### RAG - RetrievalQAWithSourcesChain using Relationship_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "880c4913",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "relationship_retriever = relationship_vector.as_retriever()\n",
    "\n",
    "from langchain.chains import RetrievalQAWithSourcesChain\n",
    "relationship_chain = RetrievalQAWithSourcesChain.from_chain_type(\n",
    "    llm, chain_type=\"stuff\", retriever=relationship_retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "f99dac89",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'The Customer Service Staff and Back Office Staff of the Big Bank plc use the Mainframe Banking System. The Internet Banking System gets account information from and makes payments using the Mainframe Banking System.\\n\\n',\n",
       " 'sources': 'Big Bank plc (multiple sources)'}"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## information found from relationship index\n",
    "relationship_chain.invoke(\n",
    "    {\"question\": \"What system is used by Service Staff\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "6b89e09c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'The document does not provide information on the system used by the Customer Service Staff.\\n\\n',\n",
       " 'sources': ''}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## information not found from element index\n",
    "chain.invoke(\n",
    "    {\"question\": \"What system is used by Service Staff\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a038facf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'A personal banking customer can view account balances and make payments using a mobile app, internet banking system, or single-page application.\\n\\n',\n",
       " 'sources': 'Big Bank plc (multiple sources)'}"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## information found from relationship index\n",
    "relationship_chain.invoke(\n",
    "    {\"question\": \"what can provide a summary of a customer's bank accounts?\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "871786b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'The \"Accounts Summary Controller\" component of the \"Big Bank plc\" software system provides a summary of a customer\\'s bank accounts.\\n\\n',\n",
       " 'sources': '1, 2, 3.'}"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## information found from element index\n",
    "chain.invoke(\n",
    "    {\"question\": \"what can provide a summary of a customer's bank accounts?\"},\n",
    "    return_only_outputs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6102ba1",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0ec64ab9",
   "metadata": {},
   "source": [
    "### RAG using Relationship_index and Element_index\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
